import baostock as bs
import pandas as pd
import numpy as np
from joblib import load
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib as mpl

# 设置中文字体
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定中文字体为黑体
mpl.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题

# 线性回归方法
def get_stock_data(code, start_date, end_date):
    # 登录证券宝
    lg = bs.login()
    if lg.error_code!= '0':
        print(f"登录失败: {lg.error_code}, {lg.error_msg}")
        exit()
    # 获取股票的历史数据
    rs = bs.query_history_k_data_plus(code,
                                 "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",
                                 start_date=start_date, end_date=end_date,
                                 frequency="d", adjustflag="3")
    if rs.error_code!= '0':
        print(f"获取数据失败，错误码：{rs.error_code}，错误信息：{rs.error_msg}")
        exit()
    # 转为 DataFrame
    data = []
    while (rs.error_code == '0') & rs.next():
        data.append(rs.get_row_data())
    df = pd.DataFrame(data, columns=rs.fields)
    # 登出证券宝
    bs.logout()
    return df

def predict_stock(code):
    # 加载模型
    model = load('model_line_0119.h5')
    # 获取2024年12月的股票数据
    df = get_stock_data(code, start_date='2024-11-01', end_date='2024-12-31')
    # 数据类型转换
    for col in ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'pctChg', 'peTTM', 'pbMRQ', 'psTTM', 'pcfNcfTTM', 'turn']:
        df[col] = pd.to_numeric(df[col])
    # 提取输入数据
    X = df[['open', 'high', 'low', 'close', 'volume', 'amount', 'pctChg']].values
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    # 预测涨幅
    y_pred = model.predict(X_scaled)
    # 放大预测数据 10 倍
    y_pred = y_pred * 1
    return y_pred

if __name__ == "__main__":
    # 用户输入股票代码
    stock_code = input("请输入股票代码（后6位数字）：")
    # 在前面加上前缀'sh.'或'sz.'
    if stock_code.startswith('6'):
        stock_code = 'sh.' + stock_code
    else:
        stock_code = 'sz.' + stock_code
    # 预测涨幅
    predicted_pct_chg = predict_stock(stock_code)
    # 获取实际涨幅
    actual_df = get_stock_data(stock_code, start_date='2024-11-01', end_date='2024-12-31')
    actual_pct_chg = actual_df['pctChg'].astype(float).values
    # 将预测涨幅添加到原始表格的最后一列
    actual_df['预测涨幅'] = predicted_pct_chg
    # 保存表格为model_out_0119.xlsx
    actual_df.to_excel('model_out_0119.xlsx', index=False)
    # 处理日期数据，只显示月份和日期
    dates = actual_df['date'].apply(lambda x: x[5:])
   
    # 绘制图表
    fig, ax = plt.subplots()
    ax.plot(dates, actual_pct_chg, label='实际涨幅')
    ax.plot(dates, predicted_pct_chg, label='预测涨幅')
    # 设置日期格式
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
    # 设置日期间隔为15个数据点
    ax.xaxis.set_major_locator(mdates.AutoDateLocator(minticks=3, maxticks=7, interval_multiples=False))
    plt.title(f"{stock_code}在2024年12月的涨幅预测")
    plt.xlabel('日期')
    plt.ylabel('涨幅')
    plt.legend()
    plt.show()